
import numpy as np
from PIL import Image

from straug.blur import GaussianBlur, DefocusBlur, MotionBlur, GlassBlur, ZoomBlur
from straug.camera import Contrast, Brightness, JpegCompression, Pixelate
from straug.geometry import Rotate, Perspective, Shrink, TranslateX, TranslateY
from straug.noise import GaussianNoise, ShotNoise, ImpulseNoise, SpeckleNoise
from straug.pattern import VGrid, HGrid, Grid, RectGrid, EllipseGrid
from straug.process import Posterize, Solarize, Invert, Equalize, AutoContrast, Sharpness, Color
from straug.warp import Curve, Distort, Stretch
from straug.weather import Fog, Snow, Frost, Rain, Shadow

class STRAug:
    """Randomly crop images and make sure to contain text instances.

    Args:
        target_size (tuple or int): (height, width)
        positive_sample_ratio (float): The probability of sampling regions
            that go through positive regions.
    """

    def __init__(self,iters=5):
        rng = np.random.default_rng()
        ops = [Curve(rng=rng), Rotate(rng=rng), Perspective(rng), Distort(rng), Stretch(rng), Shrink(rng),
               TranslateX(rng),
               TranslateY(rng), VGrid(rng), HGrid(rng), Grid(rng), RectGrid(rng), EllipseGrid(rng)]
        ops.extend([GaussianNoise(rng), ShotNoise(rng), ImpulseNoise(rng), SpeckleNoise(rng)])
        ops.extend([GaussianBlur(rng), DefocusBlur(rng), MotionBlur(rng), GlassBlur(rng), ZoomBlur(rng)])
        ops.extend([Contrast(rng), Brightness(rng),  Pixelate(rng)
                    ])
        ops.extend([Fog(rng), Snow(rng), Frost(rng), Rain(rng), Shadow(rng)])
        ops.extend(
            [Posterize(rng), Solarize(rng), Invert(rng), Equalize(rng), AutoContrast(rng), Sharpness(rng), Color(rng)])
        self.augs=ops
        self.iters=iters


    def __call__(self, results):

        img=results['img']
        h,w,c=img.shape
        # if h<32:
        #     return results
        img=Image.fromarray(img)
        cur_ops=np.random.choice(self.augs,self.iters)
        for op in cur_ops:
            mag=np.random.randint(-1,1)
            img= op(img, mag=mag)
        results['img']=np.array(img)
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        return repr_str
    
if __name__=='__main__':
    from glob import glob
    import cv2
    dirs='/media/wsl/SB@data/all_dataset/COCO-Text-words-trainval/val_words/'
    aug=STRAug(iters=2)
    files=glob(dirs+'*.*')
    for file in files:
        im=cv2.imread(file)
        ori=im.copy()
        res=aug({'img':im})
        im=res['img']
        cv2.imshow('a',im)
        cv2.imshow('ori',ori)
        cv2.waitKey(0)
